import os
from uuid import uuid4
from loguru import logger
from contextvars import ContextVar
# 使用任务request_id来实现全链路日志追踪
_trace_id: ContextVar[str] = ContextVar('x_trace_id', default="")           # 任务ID
_x_request_id: ContextVar[str] = ContextVar('x_request_id', default="")     # 请求ID

# FastAPI 日志链路追踪
class TraceID:
    """全链路追踪ID"""
    @staticmethod
    def set(req_id: str) -> ContextVar[str]:
        """设置请求ID，外部需要的时候，可以调用该方法设置
        Returns:
            ContextVar[str]: _description_
        """
        if req_id:
            req_id = uuid4().hex
        _x_request_id.set(req_id)
        return _x_request_id

    @staticmethod
    def set_trace(id: str, title: str = "task") -> ContextVar[str]:
        """设置全链路追踪ID
        Returns:
            ContextVar[str]: _description_
        """
        id = f"{title}:{id}"
        _trace_id.set(id)
        return _trace_id


def _logger_filter(record):
    record['trace_msg'] = f"{_x_request_id.get()} | {_trace_id.get()}"
    return record['trace_msg']


# 日志配置  公共参数
params = {
    "rotation": "00:00", "encoding": 'utf-8', "enqueue": True, "backtrace": True,
    "filter": _logger_filter,
    "format": "{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {trace_msg} | {name}:{function}:{line} - {message}",
}
# 去除默认控制台输出
# logger.remove()
logger.add("./logs/app-{time:YYYY-MM-DD}.log", level='INFO', retention="7 days", **params)    # info日志只保留7天

__all__ = ["logger", "TraceID"]